from abc import abstractmethod
from typing import Any, Dict, List, Optional
from continuallearning.interfaces.core.component import ComponentInterface
from .hook_interfaces import TaskAwareHookInterface
from continuallearning.interfaces.core.task_related_interface import (
    TaskIrrelevantInterface,
)


class HookManagerInterface(ComponentInterface, TaskIrrelevantInterface):
    """
    Abstract base class for hook managers.

    Hook managers coordinate multiple hooks across a model, handling
    initialization, attachment, and detachment of hooks.
    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__()
        self._initialized = False
        self._is_attached = False

    @property
    def is_initialized(self) -> bool:
        """Check if the hook manager is initialized."""
        return self._initialized

    @property
    def is_attached(self) -> bool:
        """Check if hooks are currently attached."""
        return self._is_attached

    @abstractmethod
    def _initialize_resources(self) -> None:
        """Initialize resources needed by hooks."""
        pass

    @property
    @abstractmethod
    def hooks(self) -> Dict[str, TaskAwareHookInterface]:
        """
        Get the dictionary of hooks managed by this hook manager.

        Returns:
            Dict[str, HookInterface]: Dictionary mapping hook identifiers to hook objects
        """
        pass

    def attach(self) -> "HookManagerInterface":
        """
        Initialize resources and add hooks.

        For convenience, this method combines initialization and hook addition,
        making it easier to use hook managers.

        Returns:
            Self: For method chaining
        """
        self._initialize_resources()
        self.add_hooks()
        return self

    def detach(self) -> None:
        """
        Detach and clean up all hooks.

        Ensures all hooks are properly removed from the model.
        """
        self.remove_hooks()

    @abstractmethod
    def add_hooks(self) -> None:
        """Add hooks to the model."""
        pass

    @abstractmethod
    def remove_hooks(self) -> None:
        """Remove all hooks from the model."""
        pass

    def get_num_hooks(self) -> int:
        """
        Get the number of active hooks managed by this hook manager.

        Returns:
            int: Number of active hooks
        """
        return 0

    def get_hooks(self) -> Dict[str, Any]:
        """
        Get a read-only view of the hooks managed by this hook manager.

        This method provides access to hook information for debugging and
        inspection purposes. Implementations should return a dictionary
        mapping from hook identifiers (typically module path) to hook objects.

        Returns:
            Dict[str, Any]: Dictionary of hook identifiers to hook objects
        """
        return {}

    def __enter__(self) -> "HookManagerInterface":
        """Enter context manager, attaching hooks if not already attached."""
        if not self._is_attached:
            self.attach()
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Exit context manager, detaching hooks."""
        self.detach()


class TaskAwareHookManagerInterface(HookManagerInterface):
    _current_task: int = -1

    @abstractmethod
    def register_task(
        self, task_id: int, config: Optional[Dict[str, Any]] = None
    ) -> None:
        """
        Register a new task with this hook.

        Args:
            task_id: Identifier for the new task
            config: Optional task-specific configuration
        """
        pass

    @abstractmethod
    def set_task(self, task_id: int) -> None:
        """
        Configure the model for a specific task.

        Args:
            task_id: Task identifier
        """
        self._current_task = task_id

    @abstractmethod
    def set_trainable_tasks(self, task_ids: int | List[int], **kwargs) -> None:
        """
        Prepare this hook for a specific task.

        Args:
            task_id: Identifier for the task to prepare
            **kwargs: Additional preparation parameters
        """
        pass

    @property
    @abstractmethod
    def current_task(self) -> int:
        """Get the current task ID."""
        pass

    @property
    @abstractmethod
    def registered_tasks(self) -> List[int]:
        """Get list of registered task IDs."""
        pass
